
import struct
from array import array
import mmap
import os.path
from cStringIO import StringIO

##################################################

def skip_string(x, string):
	assert x.read(len(string)) == string

def read_value(x, format):
	return struct.unpack(format, x.read(struct.calcsize(format)))[0]

def read_values(x, format):
	return struct.unpack(format, x.read(struct.calcsize(format)))

def read_uint(x):
	return struct.unpack('I', x.read(4))[0]
	#return read_value(x, 'I')

def read_ushort(x):
	return read_value(x, 'H')

def read_ulong(x):
	return read_value(x, 'Q')

def read_ubyte(x):
	return struct.unpack('B', x.read(1))[0]
	#return read_value(x, 'B')

def TVPTLG5DecompressSlide(out, stream, size, text, r):
	assert size

	flags = 0

	outi = 0

	buffer = struct.unpack('%dB'%size, stream.read(size))
	n = 0

	while n < size:
		flags >>= 1
		if not flags & 0x100:
			flags = buffer[n] | 0xff00
			n += 1

		if flags & 1:
			x0 = buffer[n]
			x1 = buffer[n+1]
			n += 2
			mpos = x0 | ((x1 & 0xf) << 8);
			mlen = (x1 & 0xf0) >> 4;
			mlen += 3
			if mlen == 18:
				mlen += buffer[n]
				n += 1
			for i in range(mlen):
				out[outi] = text[r] = text[mpos]
				outi += 1
				r += 1
				mpos += 1
				mpos &= 4096 - 1
				r &= 4096 - 1
		else:
			c = buffer[n]
			n += 1
			out[outi] = text[r] = c
			outi += 1
			r += 1
			r &= 4096 - 1

	assert n == size

	return r


def TVPTLG5ComposeColors3(out, outbuf, outi, upperi, bufi, width):
	pc0 = 0
	pc1 = 0
	pc2 = 0
	c0 = 0
	c1 = 0
	c2 = 0
	bufi0 = bufi[0]
	bufi1 = bufi[1]
	bufi2 = bufi[2]
	outbuf0 = outbuf[0]
	outbuf1 = outbuf[1]
	outbuf2 = outbuf[2]
	for x in range(width):
		c0 = outbuf0[bufi0+x]
		c1 = outbuf1[bufi1+x]
		c2 = outbuf2[bufi2+x]
		c0 += c1
		c2 += c1
		pc0 += c0
		pc1 += c1
		pc2 += c2
		out[outi+0] = (pc0 + out[upperi+0]) & 0xff
		out[outi+1] = (pc1 + out[upperi+1]) & 0xff
		out[outi+2] = (pc2 + out[upperi+2]) & 0xff
		outi += 3
		upperi += 3


def TVPTLG5ComposeColors4(out, outbuf, outi, upperi, bufi, width):
	pc0 = 0
	pc1 = 0
	pc2 = 0
	pc3 = 0
	c0 = 0
	c1 = 0
	c2 = 0
	c3 = 0
	bufi0 = bufi[0]
	bufi1 = bufi[1]
	bufi2 = bufi[2]
	bufi3 = bufi[3]
	outbuf0 = outbuf[0]
	outbuf1 = outbuf[1]
	outbuf2 = outbuf[2]
	outbuf3 = outbuf[3]
	for x in range(width):
		c0 = outbuf0[bufi0+x]
		c1 = outbuf1[bufi1+x]
		c2 = outbuf2[bufi2+x]
		c3 = outbuf3[bufi3+x]
		c0 += c1
		c2 += c1
		pc0 += c0
		pc1 += c1
		pc2 += c2
		pc3 += c3
		out[outi+0] = (pc0 + out[upperi+0]) & 0xff
		out[outi+1] = (pc1 + out[upperi+1]) & 0xff
		out[outi+2] = (pc2 + out[upperi+2]) & 0xff
		out[outi+3] = (pc3 + out[upperi+3]) & 0xff
		outi += 4
		upperi += 4

class tlg_t:
	def __init__(self, buffer, width, height, colors):
		self.buffer = buffer
		self.width = width
		self.height = height
		self.colors = colors

def read_tlg_stream(stream):

	skip_string(stream, "TLG5.0\x00raw\x1a")

	colors = read_ubyte(stream)
	assert colors in (3, 4)
	width = read_uint(stream)
	height = read_uint(stream)
	blockheight = read_uint(stream)

	blockcount = (height-1) // blockheight + 1

	# skip block size section
	stream.seek(stream.tell() + blockcount * 4)


	text = [0] * 4096
	out = [0] * width * height * colors

	outbuf = []
	for i in range(colors):
		outbuf.append( [0] * ((blockheight * width + 10 + 3) & ~3) )

	xr = 0

	prevline = None # out offset
	for y_blk in range(0, height, blockheight):
		for c in range(colors):
			#mark, size = read_values(stream, 'BI')
			mark = read_ubyte(stream)
			size = read_uint(stream)
			if not mark:
				xr = TVPTLG5DecompressSlide(outbuf[c], stream, size, text, xr)
			else:
				outbuf[c][:size] = map(ord, stream.read(size))

		y_lim = y_blk + blockheight
		if y_lim > height:
			y_lim = height

		outbufi = [0] * colors # outbufp

		for y in range(y_blk, y_lim):
			currenti = y * width * colors # out
			current_orgi = currenti

			if prevline is not None:
				if colors == 3:
					TVPTLG5ComposeColors3(out, outbuf, currenti, prevline, outbufi, width)
					outbufi[0] += width
					outbufi[1] += width
					outbufi[2] += width
				elif colors == 4:
					TVPTLG5ComposeColors4(out, outbuf, currenti, prevline, outbufi, width)
					outbufi[0] += width
					outbufi[1] += width
					outbufi[2] += width
					outbufi[3] += width
				else:
					assert False
			else:
				if colors == 3:
					pr = pg = pb = 0
					for x in range(width):
						b = outbuf[0][outbufi[0]+x]
						g = outbuf[1][outbufi[1]+x]
						r = outbuf[2][outbufi[2]+x]

						b += g
						r += g
						pb += b
						out[currenti] = pb
						currenti += 1
						pg += g
						out[currenti] = pg
						currenti += 1
						pr += r
						out[currenti] = pr
						currenti += 1
					outbufi[0] += width
					outbufi[1] += width
					outbufi[2] += width

				elif colors == 4:
					pr = pg = pb = pa = 0
					for x in range(width):
						b = outbuf[0][outbufi[0]+x]
						g = outbuf[1][outbufi[1]+x]
						r = outbuf[2][outbufi[2]+x]
						a = outbuf[3][outbufi[3]+x]

						b += g
						r += g
						pb += b
						out[currenti] = pb
						currenti += 1
						pg += g
						out[currenti] = pg
						currenti += 1
						pr += r
						out[currenti] = pr
						currenti += 1
						pa += a
						out[currenti] = pa
						currenti += 1
					outbufi[0] += width
					outbufi[1] += width
					outbufi[2] += width
					outbufi[3] += width
				else:
					assert False
			prevline = current_orgi


#	output = open('py.txt', 'w')
#	for x in range(height):
#		output.write('[')
#		for y in range(width):
#			z = (x*width + y) * colors
#			if colors == 3:
#				output.write('%s, ' % (tuple(out[z:z+3]),))
#			else:
#				output.write('%s, ' % (tuple(out[z:z+4]),))
#		output.write(']\n')

#	if colors == 4:
#		bi = 0
#		for i in range(width * height):
#			out[b+0] = out[b+0] * out[b+3] / 255 + (255 - out[b+3])
#			out[b+1] = out[b+1] * out[b+3] / 255 + (255 - out[b+3])
#			out[b+2] = out[b+2] * out[b+3] / 255 + (255 - out[b+3])
#			bi += 4
	return tlg_t(out, width, height, colors)

def read_tlg(path):
	with open(path, 'rb') as stream:
		return read_tlg_stream(stream)

def read_tlg_raw(bytes):
	stream = StringIO(bytes)
	return read_tlg_stream(stream)

def delete_tlg(tlg):
	# do nothing for C version
	pass


##################################################

def write_png(tlg, path):
	out, width, height, colors = tlg.buffer, tlg.width, tlg.height, tlg.colors
	n = width * height * colors

	import Image

	# TODO: frombuffer
	#image = Image.frombuffer('RGBA', (2, 2), [(128, 128, 128, 0), (128, 128, 128, 0), (128, 128, 128, 0), (128, 128, 128, 0),])

	if colors == 4:
		image = Image.new('RGBA', (width, height))
		image.putdata(zip(out[2:n:4], out[1:n:4], out[:n:4], out[3:n:4]))
		image.save(path, "PNG")
	elif colors == 3:
		image = Image.new('RGB', (width, height))
		image.putdata(zip(out[2:n:3], out[1:n:3], out[:n:3]))
		image.save(path, "PNG")

def tlg_file_to_png_file(tlg, png):
	t = read_tlg(tlg)
	try:
		write_png(t, png)
	finally:
		delete_tlg(t)

def tlg_buffer_to_png_file(tlg, png):
	t = read_tlg_raw(tlg)
	try:
		write_png(t, png)
	finally:
		delete_tlg(t)

##################################################

def write_bmp_stream(tlg, stream):
	header_size = 14
	dib_header_size = 40
	file_size = tlg.width * tlg.height * tlg.colors + header_size + dib_header_size

	stream.write('BM')
	stream.write(struct.pack('<I', file_size))
	stream.write('\0\0\0\0')
	stream.write(struct.pack('<I', header_size+dib_header_size))

	stream.write(struct.pack('<I', dib_header_size))
	stream.write(struct.pack('<I', tlg.width))
	stream.write(struct.pack('<i', -tlg.height)) # XXX: is it commonly supported?
	stream.write(struct.pack('<H', 1))
	if tlg.colors == 3:
		stream.write(struct.pack('<H', 24))
	elif tlg.colors == 4:
		stream.write(struct.pack('<H', 32))
	else:
		raise NotImplementedError()

#	stream.write('<I', 2835)
#	stream.write('<I', 2835)
#	stream.write('<I', 0)
#	stream.write('<I', 0)
#	stream.write('\x00\x00\xff\x00')
#	stream.write('\x00\xff\x00\x00')
#	stream.write('\xff\x00\x00\x00')
#	stream.write('\x00\x00\x00\xff')
#	stream.write('<I', 0x57696E20)

	stream.write(struct.pack('<I', 0))
	stream.write(struct.pack('<I', 0))
	stream.write(struct.pack('<I', 0x0ec4)) # XXX
	stream.write(struct.pack('<I', 0x0ec4)) # XXX
	stream.write(struct.pack('<I', 0))
	stream.write(struct.pack('<I', 0))

	for i in xrange(tlg.width*tlg.height*tlg.colors):
		stream.write(chr(tlg.buffer[i]))


def tlg_string_to_bmp_string(tlg):
	t = read_tlg_raw(tlg)
	try:
		buffer = StringIO()
		write_bmp_stream(t, buffer)
		return buffer.getvalue()
	finally:
		delete_tlg(t)

def tlg_file_to_bmp_file(tlg_path, bmp_path):
	with open(tlg_path, 'rb') as tlg:
		with open(bmp_path, 'wb') as bmp:
			bmp.write(tlg_string_to_bmp_string(tlg.read()))

##################################################

try:
	import tlg_c
	read_tlg = tlg_c.read_tlg
	delete_tlg = tlg_c.delete_tlg
	read_tlg_raw = tlg_c.read_tlg_raw
	tlg_string_to_bmp_string = tlg_c.tlg_string_to_bmp_string
	if tlg_c.write_png:
		write_png = tlg_c.write_png
		tlg_file_to_png_file = tlg_c.tlg_file_to_png_file
except:
	pass

if __name__ == '__main__':
	#import cProfile
	#cProfile.run("read_tlg('5.tlg')", 'prof_output')
	#import pstats
	#p = pstats.Stats('prof_output')
	#p.sort_stats('time')
	#p.print_stats()
	
	
	import tlg_c
	import time
	start = time.time()
	tlg_file_to_bmp_file('02.tlg', '02.bmp')
	print time.time() - start

